#!/usr/bin/python
# -*- encoding: utf-8 -*-

import torch
from libs.utils import AverageMeter, color_print
import importlib

class BaseTrainer:
    def __init__(self, userconfig, train_data_loader, test_data_loader, checkpoint_dir, sample_dir):
        self.userconfig = userconfig
        self.checkpoint_dir = checkpoint_dir
        self.sample_dir = sample_dir
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.train_data_loader = train_data_loader
        self.test_data_loader = test_data_loader
        self.train_image_len = len(self.train_data_loader)
        self.test_image_len = len(self.test_data_loader)
        print("train_dataset=%d | test_dataset=%d| train_batch_size=%d | test_batch_size=%d" % (self.train_image_len,
                                                                                                self.test_image_len,
                                                                                                self.userconfig.batch_size,
                                                                                                self.userconfig.test_batch_size))

        print("trainer=%s | net_name=%s| input_shape:%s | output_shape=%s" % (userconfig.trainer,
                                                                              userconfig.net_name,
                                                                              str(userconfig.input_shape),
                                                                              str(userconfig.output_shape)))

        print("samplestep %d | printstep %d | num_classes=%d" % (self.userconfig.sample_intervel,
                                                                 self.userconfig.log_intervel,
                                                                 self.userconfig.num_classes))
        self.build_model()
        self.resume_model()

        # train infor monitor
        self.train_loss_monitor = AverageMeter("train_loss")
        self.train_aiou_monitor = AverageMeter("train_aiou")
        self.train_biou_monitor = AverageMeter("train_biou")
        self.train_mae_monitor = AverageMeter("train_mae")
        self.train_dice_coef_monitor = AverageMeter("train_dice_coef")
        self.train_bce_monitor = AverageMeter("train_bce")
        self.train_focal_loss_monitor = AverageMeter("train_focal_loss")

        self.train_monitor_dict = {"train_loss": self.train_loss_monitor, "aiou": self.train_aiou_monitor,
                                   "biou": self.train_biou_monitor, "mae": self.train_mae_monitor,
                                   "dice_coef": self.train_dice_coef_monitor, "bce": self.train_bce_monitor,
                                   "focal_loss": self.train_focal_loss_monitor}

        # test infor monitor
        self.test_loss_monitor = AverageMeter("test_loss")
        self.test_aiou_monitor = AverageMeter("test_aiou")
        self.test_biou_monitor = AverageMeter("test_biou")
        self.test_mae_monitor = AverageMeter("test_mae")
        self.test_dice_coef_monitor = AverageMeter("test_dice_coef")
        self.test_bce_monitor = AverageMeter("test_bce")
        self.test_focal_loss_monitor = AverageMeter("test_focal_loss")

        self.test_monitor_dict = {"test_loss": self.test_loss_monitor, "aiou": self.test_aiou_monitor,
                                  "biou": self.test_biou_monitor, "mae": self.test_mae_monitor,
                                  "dice_coef": self.test_dice_coef_monitor, "bce": self.test_bce_monitor,
                                  "focal_loss": self.test_focal_loss_monitor}

    def update_train_monitor(self, infor_dict):
        for key, value in infor_dict.items():
            self.train_monitor_dict[key].update(value)

    def update_test_monitor(self, infor_dict):
        for key, value in infor_dict.items():
            self.test_monitor_dict[key].update(value)

    def infor_dict_to_str(self, infor_dict):
        output_str = ""
        for key, value in infor_dict.items():
            output_str += (" %s=%.4f |" % (key, value))
        return output_str[:-1]

    def reset_train_monitor(self):
        self.train_loss_monitor.reset()
        self.train_aiou_monitor.reset()
        self.train_biou_monitor.reset()
        self.train_mae_monitor.reset()
        self.train_dice_coef_monitor.reset()
        self.train_bce_monitor.reset()
        self.train_focal_loss_monitor.reset()

    def reset_test_monitor(self):
        self.test_loss_monitor.reset()
        self.test_aiou_monitor.reset()
        self.test_biou_monitor.reset()
        self.test_mae_monitor.reset()
        self.test_dice_coef_monitor.reset()
        self.test_bce_monitor.reset()
        self.test_focal_loss_monitor.reset()

    def print_train_monitor_infor(self):
        output_str = "train phase: "
        for key, value in self.train_monitor_dict.items():
            if value.is_use:
                output_str += ("%s=%.4f " % (value.name, value.avg))
        color_print(output_str, color="red")

    def print_test_monitor_infor(self):
        output_str = "test phase: "
        for key, value in self.test_monitor_dict.items():
            if value.is_use:
                output_str += ("%s=%.4f " % (value.name, value.avg))
        color_print(output_str, color="blue")

    def ckpt_name(self, surfix=".pth"):
        output_str = self.userconfig.net_name+"_"
        for key, value in self.test_monitor_dict.items():
            if value.is_use and key != "test_loss":
                output_str += ("%s-%.3f_" % (value.name, value.avg))

        for key, value in self.train_monitor_dict.items():
            if value.is_use and key != "train_loss":
                output_str += ("%s-%.3f_" % (value.name, value.avg))
        output_str += surfix
        return output_str

    def build_model(self):
        module = importlib.import_module("nets.%s.model" % self.userconfig.net_name)
        self.net = module.get_seg_model(self.userconfig.num_classes)
        self.net.to(self.device)

    def resume_model(self):
        if self.userconfig.model_path != '':
            self.net.load_state_dict(torch.load(self.userconfig.model_path, map_location=self.device))
            print("load snapshot %s" % self.userconfig.model_path)

    def train(self):
        pass

    def eval(self):
        pass